from typing import List
import pymongo
from bson import ObjectId
from datetime import datetime
import uuid
import pandas as pd
from model.stock import Stock
from utils.databaseConfig import databaseConfig
class StockConnector:
    def __init__(self) -> None:
        client = pymongo.MongoClient(databaseConfig.database_url)
        db = client[databaseConfig.database_name]
        self.collection = db["stock"]


    def getStockList(self) -> List[Stock]:
        stock_list = []
        for i in self.collection.find():
            if i['compare'] is not True:
                temp_stock = Stock(str(i["_id"]),i["name"],i["stock_code"])
                stock_list.append(temp_stock)
        return stock_list

    def getCompareStockList(self) -> List[Stock]:
        stock_list = []
        for i in self.collection.find():
            if i['compare'] is True:
                temp_stock = Stock(str(i["_id"]),i["name"],i["stock_code"])
                stock_list.append(temp_stock)
        return stock_list
    
    def getAttributeDataInRange(self,stock_code:str,attribute_name:str,start_date:datetime,end_date:datetime):
        # 定义开始日期和结束日期
        query = {"stock_code":stock_code}
        data = self.collection.find_one(query)["data"]
        df = pd.DataFrame(data)
        df = df[["date",attribute_name]]
        df = df[df['date'].between(start_date, end_date)]
        df.sort_values('date', inplace=True)
        df.reset_index(drop=True)
        return df
    
    def getAttributeOfStockOnDate(self,stock_code:str,date:datetime,attribute_name:str,offdays=0)->float:
        query = {"stock_code":stock_code}
        data = self.collection.find_one(query)["data"]
        df = pd.DataFrame(data)
        df = df[["date",attribute_name]]
        if offdays == 0:
            df = df[df['date'] == date]
            return df.iloc[0,1]
        else:
            index = df.loc[df['date']==date].index[0]
            values = df.loc[index+1:index+offdays,attribute_name].tolist()
            return values    
    
        
    def getStockCodeById(self,id:str)->str:
        query = {"_id":ObjectId(id)}
        stock = self.collection.find_one(query)["stock_code"]
        return stock
    
    def getStockNameByCode(self,stock_code:str)->str:
        query = {"stock_code":stock_code}
        name = self.collection.find_one(query)["name"]
        return name
    
    
stockConnector = StockConnector()
